/*
* Copyright (C) Microsoft. All rights reserved.  
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not 
* use this file except in compliance with the License.  You may obtain a copy 
* of the License at http://www.apache.org/licenses/LICENSE-2.0  
* 
* THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED 
* WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE, 
* MERCHANTABLITY OR NON-INFRINGEMENT. 
*
* See the Apache License, Version 2.0 for specific language governing 
* permissions and limitations under the License.
*
*/
// this header is not meant to be stand-alone, it should be included in callstack.hpp
#pragma once

#include <baja/common.hpp>
#include <baja/trace.hpp>
#include <baja/assert.hpp>
#include <baja/error.hpp>
#include <baja/callstack.hpp>
#include <baja/util.hpp>

#include <Windows.h>
#include <dbghelp.h>
#pragma comment(lib, "dbghelp")

namespace baja
{
struct callstack
{
    static std::string trace(int32 skipCount = 0, int32 totalCount = -1, bool noThrow = false)
    {
        // make sure we don't try to capture the stack on an error to prevent infinite recursion
        diagnostics::config::noStack = true;

        std::string s;
        withFinally::run(
        [&]() 
        {
            BAJA_CHECK_ARG(skipCount >= 0);
            
            try
            {
                callstack::init();
                s = callstack::traceInternal(skipCount, totalCount);
            }
#ifdef BAJA_CONFIG_ENABLE_TRACING
            catch (const std::exception& e)
#else
            catch (std::exception)
#endif
            {
                if (noThrow)
                {
                    BAJA_TRACE(boost::format("Exception was thrown during stack capture, stack trace won't be available...\n\tdetails: %s") % e.what());
                }
                else
                {
                    throw;
                }
            }
        },
        [&]()
        {
            callstack::cleanup();
            baja::diagnostics::config::noStack = false;
        });

        return s;
    }

private:

    static void init()
    {
        std::string searchPath = callstack::getSymbolSearchPath();

        BOOL b = SymInitialize(GetCurrentProcess(), searchPath.c_str(), TRUE);
        if (!b)
        {
            // If failed to initialize, try uninitializing and then retry
            // This happens when running with some assert libraries which apparently don't uninitialize things properly
            if (HRESULT_FROM_WIN32(GetLastError()) == E_INVALIDARG)
            {
                SymCleanup(GetCurrentProcess());
                b = SymInitialize(GetCurrentProcess(), searchPath.c_str(), TRUE);
            }
        }

        win32_result res = !!b;

        DWORD options = SymGetOptions();
        options |= SYMOPT_LOAD_LINES;
        options |= SYMOPT_FAIL_CRITICAL_ERRORS;
        options |= SYMOPT_CASE_INSENSITIVE;
        SymSetOptions(options);
    }

    static void cleanup()
    {
        SymCleanup(GetCurrentProcess());
    }

    static std::string traceInternal(int32 skipCount, int32 totalCount)
    {
        std::ostringstream sstm;

        HANDLE proc = GetCurrentProcess();
        
        CONTEXT c;
        std::memset(&c, 0, sizeof(CONTEXT));
        RtlCaptureContext(&c);

        STACKFRAME64 s;
        std::memset(&s, 0, sizeof(STACKFRAME64));

        s.AddrPC.Offset = c.Eip;
        s.AddrPC.Mode = AddrModeFlat;
        s.AddrFrame.Offset = c.Ebp;
        s.AddrFrame.Mode = AddrModeFlat;
        s.AddrStack.Offset = c.Esp;
        s.AddrStack.Mode = AddrModeFlat;

        std::vector<byte> sym;
        sym.resize(sizeof(IMAGEHLP_SYMBOL64) + 64); // TEST: make sure this still works, I used to do a new here
        IMAGEHLP_SYMBOL64* psym = reinterpret_cast<IMAGEHLP_SYMBOL64*>(sym.data());
        std::memset(psym, 0, sizeof(IMAGEHLP_SYMBOL64) + 64);
        psym->SizeOfStruct = sizeof(IMAGEHLP_SYMBOL64);
        psym->MaxNameLength = 64;

        IMAGEHLP_LINE64 line;
        std::memset(&line, 0, sizeof(IMAGEHLP_LINE64));
        line.SizeOfStruct = sizeof(IMAGEHLP_LINE64);

        IMAGEHLP_MODULE64 module;
        std::memset(&module, 0, sizeof(IMAGEHLP_MODULE64));
        module.SizeOfStruct = sizeof(IMAGEHLP_MODULE64);

        // NOTE: hack
        // skip first 6 frames since it'll be the call to callstack::trace
        skipCount += 6;

        int32 frame = 0;
        for (frame = 0; ; frame++)
        {
            if (!callstack::execute(
                "StackWalk64",
                [&]() 
                {
                    return StackWalk64(
                        IMAGE_FILE_MACHINE_I386, // NYI: X86
                        proc,
                        GetCurrentThread(),
                        &s,
                        &c,
                        nullptr,
                        SymFunctionTableAccess64,
                        SymGetModuleBase64,
                        nullptr
                        );
                })) break;

            if (s.AddrPC.Offset == s.AddrFrame.Offset)
            {
                break; // endless callstack ?
            }
            
            if (s.AddrPC.Offset != 0 && frame >= skipCount)
            {
                DWORD offset = 0;

                // Get line info
                if (!callstack::execute("SymGetLineFromAddr64", [&]() { return SymGetLineFromAddr64(proc, s.AddrPC.Offset, &offset, &line); }))
                {
                    break;
                }
                        
                // Get module info
                if (!callstack::execute("SymGetModuleInfo64", [&]() { return SymGetModuleInfo64(proc, s.AddrPC.Offset, &module); }))
                {
                    break;
                }

                // Get symbol info
                DWORD64 offset2 = 0;
                if (!callstack::execute("SymGetSymFromAddr64", [&]() { return SymGetSymFromAddr64(proc, s.AddrPC.Offset, &offset2, psym); }))
                {
                    break;
                }

                boost::format fmt;
 
                char* moduleName = module.ModuleName[0] != 0 ? module.ModuleName : "(no module)";
                if (line.FileName[0] != 0)
                {
                    // tabulation logic:
                    // 2 chars: frame number
                    // 32 chars (max): module name
                    // 64 chars (max): symbol name, as defined above 
                    // 10 chars (max): offset (for 32 bit) // NYI: X86
                    // total: 108 max, so we'll make it 64 which should be closer to average
                    fmt = boost::format("\n%02d %s!%s + 0x%x%|64t| %s, %d") % frame % moduleName % psym->Name % s.AddrPC.Offset % line.FileName % line.LineNumber;
                }
                else // Line info not available
                {
                    fmt = boost::format("\n%02d %s!%s + 0x%x%|64t| (source code not available)\n") % frame % moduleName % psym->Name % s.AddrPC.Offset;
                }

                sstm << fmt;
            }

            if ((s.AddrReturn.Offset == 0) || 
                (totalCount > 0 && frame - skipCount >= totalCount - 1))
            {
                break;
            }
        }

        return sstm.str();
    }

    BAJA_INLINE static bool execute(const char* function, std::function<BOOL(void)> f)
    {
        BOOL b = f();
        if (!b)
        {
            DWORD err = GetLastError();
            if (err == ERROR_INVALID_ADDRESS || err == ERROR_NOACCESS)
            {
                return false; // we are done
            }
            else
            {
                win32_result res;
                res.setMessage(function);
                res = false;
            }
        }

        return true;
    }

    static std::string getSymbolSearchPath()
    {
        std::string s;
        std::vector<char> v; 

        // Add current directory
        DWORD bufferSize = GetCurrentDirectoryA(0, NULL);
        if (bufferSize <= 0) BAJA_TRACE(boost::format("GetCurrentDirectoryA failed to get size: %d") % GetLastError());
        if (bufferSize > 0)
        {
            v.resize(bufferSize);
            DWORD i = GetCurrentDirectoryA(bufferSize, v.data());
            if (i <= 0) BAJA_TRACE(boost::format("GetCurrentDirectoryA failed to get data: %d") % GetLastError());
            if (i >= bufferSize) BAJA_TRACE("GetCurrentDirectoryA: buffer not large enough");
    
            if (i > 0 && i < bufferSize)
            {
                s.append(v.data());
                s.append(";");
            }
        }

        static const char* vars[] = { "_NT_SYMBOL_PATH", "_NT_ALTERNATE_SYMBOL_PATH" };

        // Add environment variables
        std::for_each(
            std::begin(vars), 
            std::end(vars),
            [&v, &s](const char*& var)
        {
            DWORD bufferSize = GetEnvironmentVariableA(var, NULL, 0);
            if (bufferSize <= 0 && GetLastError() != ERROR_ENVVAR_NOT_FOUND)
                BAJA_TRACE(boost::format("GetEnvironmentVariableA failed: %d") % GetLastError());

            if (bufferSize > 0)
            {
                v.resize(bufferSize);
                DWORD i = GetEnvironmentVariableA(var, v.data(), bufferSize);
                if (i == 0) BAJA_TRACE(boost::format("GetEnvironmentVariableA failed: %d") % GetLastError());
                if (i >= bufferSize) BAJA_TRACE("GetEnvironmentVariableA: buffer not large enough");
                if (i > 0 && i < bufferSize)
                {
                    s.append(v.data());
                    s.append(";");
                }
            }
        });

        BAJA_TRACE(boost::format("Using symbol search path: %s") % s);

        return s;
    }

private:
    int32 _frameCount;
    bool _noThrow;
    bool _init;
};

} // namespace baja